from tensorflow.keras.datasets import mnist
from tensorflow.keras.utils import to_categorical

from tensorflow.keras import layers
from tensorflow.keras.models import Model
from tensorflow.keras import optimizers
from tensorflow.keras.losses import categorical_crossentropy








inputs = layers.Input((28,28,1))
x = layers.Conv2D(64,3)(inputs)
x = layers.BatchNormalization(center=True,scale=False)(x)
x = layers.Activation('relu')(x)
x = layers.Conv2D(64,3,strides=2)(x)
x = layers.BatchNormalization(center=True,scale=False)(x)
x = layers.Activation("relu")(x)
x = layers.Conv2D(128,5)(x)
x = layers.BatchNormalization(center=True,scale=False)(x)
x = layers.Activation('relu')(x)
x = layers.Conv2D(128,5)(x)
x = layers.BatchNormalization(center=True,scale=False)(x)
x = layers.Activation("relu")(x)
x = layers.Flatten()(x)
x = layers.Dense(100)(x)
x = layers.BatchNormalization(center=True,scale=False)(x)
x = layers.Activation('relu')(x)
x = layers.Dropout(0.3)(x)
x = layers.Dense(10,activation='softmax')(x)
model = Model(inputs,x)
model.compile(optimizer=optimizers.SGD(momentum=0.8,nesterov=True),loss=categorical_crossentropy,metrics=['accuracy'])
model.summary()



(x_train,y_train),(x_test,y_test) = mnist.load_data()
X_train = x_train.reshape((60000,28,28,1))/255.
X_test = x_test.reshape((10000,28,28,1))/255.

y_train = to_categorical(y_train,num_classes=10)
y_test = to_categorical(y_test,num_classes=10)

model.fit(X_train,y_train,batch_size=128,epochs=5,validation_split=0.2,verbose=2)

#model.fit(X_train,y_train,batch_size=128,epochs=18,verbose=2)
print(model.evaluate(X_test,y_test))

model.save("teacher_model/")
